# ai_client.py (异步版本 - 专用豆包)
import aiohttp
from openai import AsyncOpenAI, APIStatusError
import json
import re
import logging
from config import Config

class AIClientAsync:
    def __init__(self, config: Config, task_type: str):
        self.config = config
        self.task_type = task_type
        self.client = self._create_client()

    def _create_client(self):
        """根据 task_type 创建不同的豆包客户端"""
        if self.task_type == "generate":
            return DoubaoAIClient(self.config, model=self.config.DOUBAO_GENERATE_MODEL)
        elif self.task_type == "complete":
            return DoubaoAIClient(self.config, model=self.config.DOUBAO_COMPLETE_MODEL)
        elif self.task_type == "analyze_elements": # 新增分支
            return DoubaoAIClient(self.config, model=self.config.DOUBAO_ANALYZE_ELEMENTS_MODEL, task_type=self.task_type)
        else:
            valid_task_types = ["generate", "complete", "analyze_elements"]
            raise ValueError(f"task_type 必须是以下之一: {', '.join(valid_task_types)}")

    async def chat(self, messages, name_length=None):
        """委托给具体的 AI 客户端"""
        return await self.client.chat(messages, name_length)


class DoubaoAIClient:  # 统一的豆包客户端
    def __init__(self, config: Config, model: str, task_type: str):
        self.config = config
        self.model = model  # 直接传入模型名称
        self.task_type = task_type
        self.async_client = AsyncOpenAI(
            api_key=config.DOUBAO_API_KEY,
            base_url=config.DOUBAO_BASE_URI,
        )

    async def chat(self, messages, name_length=None):
        if name_length is not None and name_length not in self.config.DOUBAO_NAME_LENGTHS:
            logging.warning(f"豆包模型不支持生成 {name_length} 个字的名字。")
            return None
        try:
            async with aiohttp.ClientSession() as session:
                response = await self.async_client.chat.completions.create(
                    model=self.model,  # 使用传入的模型
                    messages=messages,
                )
                if response.choices:
                    content = response.choices[0].message.content
                    return self._parse_content(content)
                else:
                    return None

        except APIStatusError as e:
            print(f"----- 豆包 AI API 错误 -----")
            print(f"状态码: {e.status_code}")
            print(f"错误类型: {e.type}")
            print(f"错误消息: {e.message}")
            print(f"错误详情: {e.response}")
            print("----------------------------")
            return None

        except Exception as e:
            logging.error(f"豆包 AI 请求失败: {e}")
            return None

    def _parse_content(self, content):
        """尝试从返回内容中提取 JSON 数据"""
        try:
            return json.loads(content)
        except json.JSONDecodeError:
            pass

        match = re.search(r"`(\[.*?\])`", content, re.DOTALL)
        if match:
            try:
                return json.loads(match.group(1))
            except json.JSONDecodeError:
                pass

        return None